package com.gitee.wsl.ext.map

import com.gitee.wsl.exception.DuplicateKeyException
import com.gitee.wsl.ext.collection.assocWith
import com.gitee.wsl.ext.collection.union


/**
 * Merges two maps together using the given [reduce] function.
 * By default, the reduce function keeps the value from [this] map.
 */
inline fun <K, V> Map<K, V>.mergeReduce(other: Map<K, V>, reduce: (V, V) -> V = { a, _ -> a }): Map<K, V> {
    return mergeReduceTo(LinkedHashMap(this.size + other.size), other, reduce)
}

/**
 * Merges two maps together using the given [reduce] function into the given [destination].
 * By default, the reduce function keeps the value from [this] map.
 */
inline fun <K, V, M : MutableMap<K, V>> Map<K, V>.mergeReduceTo(destination: M, other: Map<K, V>, reduce: (V, V) -> V = { a, _ -> a }): M {
    destination.putAll(this)

    other.forEach { (key, value) ->
        destination[key] = destination[key]?.let { reduce(it, value) } ?: value
    }
    return destination
}


/**
 * For each key, merges all the values into one common list.
 */
fun <K, V : Any> Iterable<Map<K, V>>.merge(): Map<K, List<V>> {
    val keys = this.map { it.keys }.union()
    return keys.assocWith { key -> this.mapNotNull { map -> map[key] } }
}

/**
 * For each key, merges all the list into one common list.
 */
fun <K, V : Any> Iterable<Map<K, List<V>>>.flatMerge(): Map<K, List<V>> {
    val keys = this.map { it.keys }.union()
    return keys.assocWith { key -> this.flatMap { map -> map[key] ?: emptyList() } }
}

infix fun <K, V> Map<K, V>.merge(another: Map<K, V>): Map<K, V> {
    val result = mutableMapOf<K, V>()
    this.toMap(result)
    another.forEach {
        if (result.containsKey(it.key)) {
            throw DuplicateKeyException("Duplicate key while merging maps")
        }
        result[it.key] = it.value
    }
    return result
}


infix fun <K, V> Map<K, V>.mergeIfNotExists(another: Map<K, V>): Map<K, V> {
    val result = mutableMapOf<K, V>()
    this.toMap(result)
    another.forEach {
        if (!result.containsKey(it.key)) {
            result[it.key] = it.value
        }
    }
    return result
}
